import time
import os
import paddle
import paddle.fluid as fluid
from modelV3 import YOLOv3
from getData import multithread_loader
import numpy as np

# 可以调整
ANCHORS = [10, 13, 16, 30, 33, 23, 30, 61, 62, 45, 59, 119, 116, 90, 156, 198, 373, 326]
# ？
ANCHOR_MASKS = [[6, 7, 8], [3, 4, 5], [0, 1, 2]]
# 超参数 阈值
IGNORE_THRESH = .7
# 昆虫种类，不能改
NUM_CLASSES = 7

# 预置学习力
def get_lr(base_lr = 0.001, lr_decay = 0.1):
    bd = [10000, 20000]
    lr = [base_lr, base_lr * lr_decay, base_lr * lr_decay * lr_decay]
    learning_rate = fluid.layers.piecewise_decay(boundaries=bd, values=lr)
    return learning_rate

if __name__ == '__main__':

    TRAINDIR = '/home/aistudio/work/insects/train'
    TESTDIR = '/home/aistudio/work/insects/test'
    VALIDDIR = '/home/aistudio/work/insects/val'

    with fluid.dygraph.guard():
        # 使用YOLOV3模型 
        # @ 模型类别数字7 ，是训练集
        model = YOLOv3(num_classes = NUM_CLASSES, is_train=True)  #创建模型
        # 学习率
        learning_rate = get_lr()
        # # 定义优化器
        opt = fluid.optimizer.Momentum(
                     learning_rate=learning_rate,
                     momentum=0.9,
                     regularization=fluid.regularizer.L2Decay(0.0005),
                     parameter_list=model.parameters())  
        #创建优化器
       
        # 定义数据读取器，训练数据读取器和验证数据读取器
        # 训练数据
        train_loader = multithread_loader(TRAINDIR, batch_size= 10, mode='train')  #创建训练数据读取器
        # 验证数据
        valid_loader = multithread_loader(VALIDDIR, batch_size= 10, mode='valid')  #创建验证数据读取器
        # 训练次数
        MAX_EPOCH = 10
        # 循环开始训练
        for epoch in range(MAX_EPOCH):
            # 取出数据
            for i, data in enumerate(train_loader()):
                # 数据分为img 图片，gt_boxes盒子，gt_labels标签分类，img_scale 图像尺寸
                img, gt_boxes, gt_labels, img_scale = data
                # 转换数据类型
                gt_scores = np.ones(gt_labels.shape).astype('float32')
                gt_scores = fluid.dygraph.to_variable(gt_scores)
                img = fluid.dygraph.to_variable(img)
                gt_boxes = fluid.dygraph.to_variable(gt_boxes)
                gt_labels = fluid.dygraph.to_variable(gt_labels)
                
                # ---------------------------------------------
                # 将图片喂给模型
                outputs = model(img)  #前向传播，输出[P0, P1, P2]
                # 获得损失函数
                loss = model.get_loss(outputs, gt_boxes, gt_labels, gtscore=gt_scores,
                                      anchors = ANCHORS,
                                      anchor_masks = ANCHOR_MASKS,
                                      ignore_thresh=IGNORE_THRESH,
                                      use_label_smooth=False)        # 计算损失函数

                # 反向传播计算梯度
                loss.backward()     
                # 更新参数
                opt.minimize(loss)  
                # 模型清理梯度
                model.clear_gradients()
                # -------------------------------------------------
                if i % 1 == 0:
                    timestring = time.strftime("%Y-%m-%d %H:%M:%S",time.localtime(time.time()))
                    print('{}[TRAIN]epoch {}, iter {}, output loss: {}'.format(timestring, epoch, i, loss.numpy()))

            # save params of model
            if (epoch % 5 == 0) or (epoch == MAX_EPOCH -1):
                fluid.save_dygraph(model.state_dict(), 'yolo_epoch{}'.format(epoch))
                
            # 每个epoch结束之后在验证集上进行测试
            model.eval()
            for i, data in enumerate(valid_loader()):
                # 数据分为img 图片，gt_boxes盒子，gt_labels标签分类，img_scale 图像尺寸
                img, gt_boxes, gt_labels, img_scale = data
                # 转换数据类型
                gt_scores = np.ones(gt_labels.shape).astype('float32')
                gt_scores = fluid.dygraph.to_variable(gt_scores)
                img = fluid.dygraph.to_variable(img)
                gt_boxes = fluid.dygraph.to_variable(gt_boxes)
                gt_labels = fluid.dygraph.to_variable(gt_labels)
                
                # 将图片喂给模型
                outputs = model(img)
                # 获得损失函数
                loss = model.get_loss(outputs, gt_boxes, gt_labels, gtscore=gt_scores,
                                      anchors = ANCHORS,
                                      anchor_masks = ANCHOR_MASKS,
                                      ignore_thresh=IGNORE_THRESH,
                                      use_label_smooth=False)
                if i % 1 == 0:
                    timestring = time.strftime("%Y-%m-%d %H:%M:%S",time.localtime(time.time()))
                    print('{}[VALID]epoch {}, iter {}, output loss: {}'.format(timestring, epoch, i, loss.numpy()))
            model.train()